#include "kernel_operator.h"

constexpr uint32_t FIXED_BATCH = 8;
constexpr uint32_t FIXED_ROWS  = 1024;
constexpr uint32_t FIXED_COLS  = 1024;
constexpr int32_t  FIXED_K     = 0;

class KernelEye {
public:
    __aicore__ inline void Init(GM_ADDR out,
                                uint32_t batch,
                                uint32_t rows,
                                uint32_t cols,
                                int32_t k)
    {
        b = batch; r = rows; c = cols; dk = k;
        blk = AscendC::GetBlockIdx();
        if (blk >= b) { valid = false; return; }
        perMatrixElems = (uint64_t)r * c;
        baseOffset = (uint64_t)blk * perMatrixElems;
        outGm.SetGlobalBuffer((__gm__ float*)out, (uint64_t)b * perMatrixElems);
        alignNum = 32 / sizeof(float); // 8
    }

    __aicore__ inline void Process()
    {
        if (!valid) return;
        for (uint32_t row = 0; row < r; ++row) {
            int32_t diagCol = (int32_t)row + dk;
            if (diagCol < 0 || (uint32_t)diagCol >= c) continue;
            uint64_t pos  = baseOffset + (uint64_t)row * c + (uint32_t)diagCol;
            uint64_t base = (pos / alignNum) * alignNum;
            uint32_t lane = (uint32_t)(pos - base);
            // 设置该向量内 lane 位置为 1
            outGm[base].SetValue(lane, 1.0f);
        }
    }

private:
    AscendC::GlobalTensor<float> outGm;
    uint32_t b=0,r=0,c=0;
    int32_t  dk=0;
    uint32_t blk=0;
    bool valid=true;
    uint64_t perMatrixElems=0;
    uint64_t baseOffset=0;
    uint32_t alignNum=0;
};

#ifdef ASCENDC_CPU_DEBUG
extern "C" __global__ __aicore__
void eye_custom(GM_ADDR out,
                uint32_t batch,
                uint32_t rows,
                uint32_t cols,
                int32_t k)
{
    KernelEye op;
    op.Init(out, batch, rows, cols, k);
    op.Process();
}
#else
extern "C" __global__ __aicore__
void eye_custom(GM_ADDR out,
                uint32_t batch,
                uint32_t rows,
                uint32_t cols,
                int32_t k)
{
    KernelEye op;
    op.Init(out, batch, rows, cols, k);
    op.Process();
}

void eye_custom_do(uint32_t blockDim,
                   void *stream,
                   uint8_t *out,
                   uint32_t batch,
                   uint32_t rows,
                   uint32_t cols,
                   int32_t k)
{
    eye_custom<<<blockDim, nullptr, stream>>>(out, batch, rows, cols, k);
}
#endif